import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.metrics import accuracy_score
import joblib

def get_data_no_index(code):
    stock_data = pd.read_csv('d:\\csv_shares_day\\' + str(code) + '-19900101-20211231.csv', encoding="utf-8")
    stock_data.columns = ["Date", "code", "name", "Close", "High", "Low", "Open", "pre_close", "subval",
                          "subrate",
                          "turnrate",
                          "Volume", "Money", "all_earn", "exchage_earn"]
    stock_data['trade_date'] = pd.to_datetime(stock_data['Date'])
    stock_data.set_index('trade_date', inplace=True)
    data_daily = stock_data.sort_index(ascending=True)
    newdata = data_daily[["Date","Close","High","Low","Open","Volume"]]
    return newdata

code = "000001"
data = get_data_no_index(code)
print(data)
# 数据准备和处理
data['Close'] = data['Close'].astype(float)
data['Open'] = data['Open'].astype(float)
data['High'] = data['High'].astype(float)
data['Low'] = data['Low'].astype(float)

# 创建标签列
data['Label'] = data['Close'].diff().gt(0).astype(int)

# 提取特征和目标变量
X = data[['Volume', 'Open', 'High', 'Low']]
y = data['Label']

# 划分训练集测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)

# 构建Pipeline
pipe = Pipeline([
    ('imputer', SimpleImputer(strategy='mean')),
    ('scaler', StandardScaler()),
    ('model', LogisticRegression())
])

# 模型训练
pipe.fit(X_train, y_train)

# 保存模型
joblib.dump(pipe, 'model.pkl')

# 测试集预测
y_pred = pipe.predict(X_test)

# 准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"准确率: {accuracy}")

import pandas as pd
import joblib

# 加载模型
loaded_model = joblib.load('model.pkl')

from services.DayKlineService import *
kline = DayKlineService()

pd.set_option('display.max_columns', None) # 展示所有列

alldate = kline.getAllData("SH603037")
new_data = pd.DataFrame(alldate)
new_data = new_data.iloc[-2:-1,:]
print(new_data)

# 新数据准备
new_data['Close'] = new_data['close'].astype(float)
new_data['Open'] = new_data['open'].astype(float)
new_data['High'] = new_data['high'].astype(float)
new_data['Low'] = new_data['low'].astype(float)
new_data['Volume'] = new_data['volume'].astype(float)
new_data['Date'] = new_data['trade_date']
# 删除Close和Date特征列
new_data.drop('Close', axis=1, inplace=True)
new_data.drop('Date', axis=1, inplace=True)

query_data = new_data[['Volume', 'Open', 'High', 'Low']]

# 预测结果
predicted_labels = loaded_model.predict(query_data)

# 输出预测结果
for i, label in enumerate(predicted_labels):
    print(f"样本{i + 1}的预测结果：{label}")